package com.tianji.aigc.service.impl;

import com.tianji.aigc.dto.ChatDTO;
import com.tianji.aigc.entity.ChatSession;
import com.tianji.aigc.enums.ChatEventTypeEnum;
import com.tianji.aigc.service.IChatService;
import com.tianji.aigc.service.IChatSessionService;
import com.tianji.aigc.vo.ChatEventVO;
import com.tianji.common.exceptions.BizIllegalException;
import com.tianji.common.utils.UserContext;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * TODO
 *
 * @Author mr.xing
 * @Date 2025/9/2 17:08
 */
@Service
public class ChatServiceImpl implements IChatService {

    @Autowired
    private ChatClient chatClient;

    @Autowired
    private IChatSessionService chatSessionService;

    //存储大模型的生成状态，这里采用ConcurrentHashMap是确保线程安全，如果分布式环境可考虑用redis
    private static Map<String,Boolean> SESSION_MAP = new ConcurrentHashMap<String,Boolean>();

    /**
     * 流式对话
     *
     * @param dto
     * @return
     */
    @Override
    public Flux<ChatEventVO> chat(ChatDTO dto) { // 重写父类或接口方法,定义公共方法chat
         //  调用service的Lambda查询功能,根据dto.get..查询对应的ChatSession对象
        ChatSession one = chatSessionService.lambdaQuery().eq(ChatSession::getSessionId, dto.getSessionId()).one();
         // 如果为空,抛出异常
        if (one==null){
            throw new BizIllegalException("对话不存在");
        }
        String converationId = UserContext.getUser()+":"+dto.getSessionId();
        // 使用之前配置的chatClient构建对话请求, prompt方法用于开始一个新的对话提示
        return chatClient.prompt()
                // 设置用户的消息内容,从dto.getQuestion获取用户提出的问题
                .user(dto.getQuestion())
                .advisors(x->x.param(AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY, converationId))
                // 使用流式方式发送请求,这样可以逐步接收响应而不是等待完整响应
                .stream()
                // 获取响应的内容部分
                .content()
                .doFirst(()-> SESSION_MAP.put(dto.getSessionId(),true))//标记开始输出
                .doOnComplete(()->SESSION_MAP.remove(dto.getSessionId()))//标记输出完成，清除标记
                .doOnError(throwable-> SESSION_MAP.remove(dto.getSessionId()))//输出异常，清除标记
                .takeWhile(takewile -> SESSION_MAP.getOrDefault(dto.getSessionId(),false))//判断是否继续输出，false停止输出，true继续输出
                 // 使用map操作符将每个内容片段x转换为chatEventVO对象,设置时间数据内容为x,类型为DATA
                .map(x->ChatEventVO.builder().eventType(ChatEventTypeEnum.DATA.getValue()).build())
                 //  使用concatWith将一个结束标记连接到流的末尾,创建一个ChatEventVo对象,类型为STOP
                .concatWith(Flux.just(ChatEventVO.builder()
                        .eventType(ChatEventTypeEnum.STOP.getValue())
                        .build()));
    }

    /**
     * 停止生成
     *
     * @param sessionId
     */
    @Override
    public void stop(String sessionId) {
        SESSION_MAP.put(sessionId,false);
    }
}
